import math
import random

def assign_cluster(x, c):
    min_dist = float('inf')
    cluster_idx = 0

    for i, center in enumerate(c):
        dist_sq = 0.0
        for xi, ci in zip(x, center):
            dist_sq += (xi - ci) ** 2
        dist = math.sqrt(dist_sq)

        if dist < min_dist:
            min_dist = dist
            max_dist = dist
            cluster_idx = i

    return cluster_idx


def Kmeans(data, k, epsilon=1e-3, iteration=100):
    if len(data) == 0:
        raise ValueError("数据集不能为空")
    if k < 1 or k > len(data):
        raise ValueError(f"k值必须在1到{len(data)}之间")
    dim = len(data[0])
    for x in data:
        if len(x) != dim:
            raise ValueError("所有数据点必须具有相同的维度")

    data_indices = list(range(len(data)))
    random.shuffle(data_indices)
    centers = [data[i].copy() for i in data_indices[:k]]

    iter_count = 0
    while iter_count < iteration:
        clusters = [[] for _ in range(k)]
        for x in data:
            cluster_idx = assign_cluster(x, centers)
            clusters[cluster_idx].append(x)

        new_centers = []
        for cluster in clusters:
            if not cluster:
                new_center = random.choice(data).copy()
            else:
                new_center = []
                for d in range(dim):
                    avg = sum(x[d] for x in cluster) / len(cluster)
                    new_center.append(avg)
            new_centers.append(new_center)

        converge = True
        for old_c, new_c in zip(centers, new_centers):
            dist_sq = 0.0
            for oc, nc in zip(old_c, new_c):
                dist_sq += (oc - nc) ** 2
            dist = math.sqrt(dist_sq)
            if dist > epsilon:
                converge = False
                break

        centers = new_centers
        iter_count += 1

        if converge:
            print(f"迭代{iter_count}次后收敛")
            break

    if not converge:
        print(f"达到最大迭代次数{iteration}，未完全收敛")

    return clusters, centers, iter_count


if __name__ == "__main__":
    random.seed(42)
    test_data = []
    for _ in range(10):
        x = 2 + random.gauss(0, 0.5)
        y = 3 + random.gauss(0, 0.5)
        test_data.append([x, y])
    for _ in range(10):
        x = 8 + random.gauss(0, 0.5)
        y = 7 + random.gauss(0, 0.5)
        test_data.append([x, y])
    for _ in range(10):
        x = 5 + random.gauss(0, 0.5)
        y = 10 + random.gauss(0, 0.5)
        test_data.append([x, y])

    k = 3
    epsilon = 1e-4
    max_iter = 50
    clusters, final_centers, iter_num = Kmeans(test_data, k, epsilon, max_iter)

    print(f"\n最终聚类中心：")
    for i, center in enumerate(final_centers):
        print(f"聚类{i + 1}中心：{[round(c, 4) for c in center]}")

    print(f"\n每个聚类的数据点数量：")
    for i, cluster in enumerate(clusters):
        print(f"聚类{i + 1}：{len(cluster)}个点")